import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import matplotlib.dates as mdates  # 必须添加此导入
import os
from utils import ensure_dir, get_font, _preprocess_dates, _set_dynamic_axis, _ensure_dir, _validate_columns, record_time
import json

# avail_font = get_font()
# assert avail_font is not None, "font error"
# 设置全局中文字体（Windows推荐微软雅黑，Linux/macOS需安装对应字体）
# plt.rcParams['font.sans-serif'] = [avail_font]  # 或 ['SimHei']（黑体）

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示为方块的问题[1,3,5](@ref)

@record_time("可视化-注册数和活跃率的组合图")
def visualize_reg_act(ddf, savedir):
    import matplotlib.dates as mdates
    
    # 1. 修复 SettingWithCopyWarning
    ddf['registration_date'] = pd.to_datetime(ddf['registration_date'], errors='coerce')
    valid_dates = ddf['registration_date'].between('2019-01-01', '2023-12-31')
    ddf = ddf[valid_dates].copy()

    # 2. 生成准确年月列
    ddf['year_month'] = ddf['registration_date'].dt.to_period('M').dt.to_timestamp()

    # 3. 统计聚合
    reg_trend = ddf.groupby('year_month').size().reset_index(name='counts')
    active_trend = ddf.groupby('year_month')['is_active'].mean().reset_index(name='active_rate')

    # 4. 创建图表
    fig, ax1 = plt.subplots(figsize=(14, 7))
    ax1.plot(reg_trend['year_month'], reg_trend['counts'], color='blue', marker='o', label='注册量')

    # 5. 添加活跃率折线图
    ax2 = ax1.twinx()
    ax2.plot(active_trend['year_month'], active_trend['active_rate']*100, 
            color='red', linestyle='--', marker='x', label='活跃率')

    # 6. 动态设置日期轴（保持原逻辑）
    date_span = (reg_trend['year_month'].max() - reg_trend['year_month'].min()).days
    if date_span > 365:
        ax1.xaxis.set_major_locator(mdates.MonthLocator(interval=3))
        def quarter_formatter(x, pos=None):
            date = mdates.num2date(x)
            quarter = (date.month - 1) // 3 + 1
            return f"{date.year}-Q{quarter}"
        ax1.xaxis.set_major_formatter(quarter_formatter)
    else:
        ax1.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))

    # 7. 图例修复关键点[1,7](@ref)
    ax1.legend(loc='upper left', bbox_to_anchor=(0, 1.02))  # 左上外侧
    ax2.legend(loc='upper right', bbox_to_anchor=(1, 1.02)) # 右上外侧
    plt.xticks(rotation=45, ha='right', fontsize=10)
    plt.tight_layout()
    plt.savefig(ensure_dir(os.path.join(savedir, 'visualize_reg_act.jpg')))
    try:
        plt.show()
    except:
        pass
    
    plt.close

@record_time("可视化-年龄和收入分布的组合图")
def visualize_age_gender_income(df, savedir):
    """
    第一个可视化函数：年龄和收入分布的组合图
    支持 pandas DataFrame 和 Dask DataFrame
    """
    # 创建子图
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
    
    # 检查是否为Dask DataFrame
    is_dask = hasattr(df, 'compute')
    
    # 1. 年龄分布直方图
    age_data = df['age'].compute() if is_dask else df['age']
    ax1.hist(age_data, bins=50, edgecolor='black')
    ax1.set_title('用户年龄分布')
    ax1.set_xlabel('年龄')
    ax1.set_ylabel('用户数量')
    
    # 2. 收入箱线图（按性别分组）
    if is_dask:
        gender_income = df.groupby('gender')['income'].compute()
    else:
        gender_income = df.groupby('gender')['income']
        
    ax2.boxplot([gender_income.get_group(g) for g in gender_income.groups],
                labels=gender_income.groups.keys())
    ax2.set_title('不同性别用户收入分布')
    ax2.set_xlabel('性别')
    ax2.set_ylabel('收入')
    
    plt.tight_layout()
    plt.savefig(ensure_dir(os.path.join(savedir, 'visualize_age_gender_income.jpg')))
    try:
        plt.show()
    except:
        pass
    
    plt.close()


@record_time("可视化-设备使用分布饼图")
def visualize_device_usage(df, savedir):
    """
    第二个可视化函数：设备使用分布饼图
    """
    def get_devices(login_history):
        try:
            if isinstance(login_history, str):
                data = json.loads(login_history)
                devices = data.get('devices', [])
                if isinstance(devices, list):
                    return devices
        except Exception as e:
            print(f"Error parsing login_history: {e}")
            print(f"Problematic login_history: {login_history[:100]}...")  # 打印前100个字符用于调试
        return []
    
    # 检查是否为Dask DataFrame
    is_dask = hasattr(df, 'compute')
    
    try:
        # 打印一些示例数据用于调试
        sample_data = df['login_history'].head(1)
        print("Sample login_history data:", sample_data.iloc[0] if not is_dask else sample_data.compute().iloc[0])
        
        # 计算设备使用分布
        devices_series = df['login_history'].apply(get_devices)
        if is_dask:
            devices_series = devices_series.compute()
        
        # 展开设备列表并创建计数
        all_devices = []
        for devices in devices_series:
            if devices:  # 确保不是空列表
                all_devices.extend(devices)
        
        if not all_devices:
            print("Warning: No valid device data found!")
            return
            
        device_counts = pd.Series(all_devices).value_counts()
        
        # 创建饼图
        plt.figure(figsize=(10, 8))
        plt.pie(device_counts.values, 
                labels=device_counts.index,
                autopct='%1.1f%%',
                startangle=90)
        plt.title(f'用户设备使用分布 (总计: {len(all_devices)})')
        
    except Exception as e:
        print(f"Error in visualization: {e}")
        return
    
    finally:
        plt.savefig(ensure_dir(os.path.join(savedir, 'visualize_device_usage.jpg')))
        try:
            plt.show()
        except:
            pass
        
        plt.close()


@record_time("可视化-商品类别分布条形图")
def visualize_purchase_categories(df, savedir):
    """
    第三个可视化函数：商品类别分布条形图
    """
    def get_categories(purchase_history):
        try:
            if isinstance(purchase_history, str):
                data = json.loads(purchase_history)
                categories = data.get('categories')
                if categories:  # 确保类别不为空
                    return categories
        except Exception as e:
            print(f"Error parsing purchase_history: {e}")
            print(f"Problematic purchase_history: {purchase_history[:100]}...")  # 打印前100个字符用于调试
        return None
    
    # 检查是否为Dask DataFrame
    is_dask = hasattr(df, 'compute')
    
    try:
        # 打印一些示例数据用于调试
        sample_data = df['purchase_history'].head(1)
        print("Sample purchase_history data:", sample_data.iloc[0] if not is_dask else sample_data.compute().iloc[0])
        
        # 计算商品类别分布
        categories_series = df['purchase_history'].apply(get_categories)
        if is_dask:
            categories_series = categories_series.compute()
        
        # 移除None值并计数
        valid_categories = categories_series.dropna()
        
        if valid_categories.empty:
            print("Warning: No valid category data found!")
            return
            
        category_counts = valid_categories.value_counts()
        
        # 创建条形图
        plt.figure(figsize=(12, 6))
        plt.bar(range(len(category_counts)), category_counts.values)
        plt.xticks(range(len(category_counts)), category_counts.index, 
                   rotation=45, ha='right')
        plt.title(f'商品类别分布 (总计: {len(valid_categories)})')
        plt.xlabel('商品类别')
        plt.ylabel('购买次数')
        plt.tight_layout()
        
    except Exception as e:
        print(f"Error in visualization: {e}")
        return
    finally:
        plt.savefig(ensure_dir(os.path.join(savedir, 'visualize_purchase_categories.jpg')))
        try:
            plt.show()
        except:
            pass
        
        plt.close()